package com.ag.improvedjava.iterable;

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Stack;

import com.ag.improvedjava.iterable.CompositeKey.KeyComparator;
import com.ag.improvedjava.iterable.Iterables.ImprovedIterableImpl;
import com.ag.improvedjava.utils.Function;
import com.ag.improvedjava.utils.OrderedIterable;
import com.ag.improvedjava.utils.Requires;
import com.ag.improvedjava.utils.comparators.Comparators;
import com.ag.improvedjava.utils.comparators.ReverseComparator;

final class OrderedIterableImpl<T, CK /* CompositeKey */> extends ImprovedIterableImpl<T> implements OrderedIterable<T> {

	private final Function<T, CK> compositeSelector; 
    private final Comparator<CK> compositeComparer;
    
    OrderedIterableImpl(Iterable<T> source, Function<T, CK> compositeSelector, Comparator<CK> compositeComparer) {
		super(source);
		this.compositeSelector = compositeSelector;
		this.compositeComparer = compositeComparer;
	}

	@Override
	public <K> OrderedIterable<T> createOrderedEnumerable(
			final Function<T, K> keySelector, Comparator<K> comparer,
			boolean descending) {
		Requires.requireNonNull(keySelector, "The key selector can not be null.");
		Requires.requireNonNull(comparer, "The comparer can not be null.");
		
		if (descending) {
			comparer = new ReverseComparator<>(comparer);
		}
		
		final Function<T, CK> primarySelector = compositeSelector;
		
		Function<T, CompositeKey<CK, K>> newKeySelector = new Function<T, CompositeKey<CK,K>>() {
			public CompositeKey<CK, K> execute(T t) {
				return new CompositeKey<>(primarySelector.execute(t), keySelector.execute(t));
			}
		};
		
		KeyComparator<CK, K> newKeyComparer = new CompositeKey.KeyComparator<>(compositeComparer, comparer);
		
		return new OrderedIterableImpl<>(iterable, newKeySelector, newKeyComparer);
	}
	
	@Override
	public Iterator<T> iterator() {
		Object[] data = Iterables.convert(iterable).toArray();
		return new OrderedIterator<>(data, compositeSelector, compositeComparer);
	}

	@Override
	public <K extends Comparable<? super K>> OrderedIterable<T> thenBy(Function<T, K> keySelector) {
		return thenBy(keySelector, Comparators.<K>Natural());
	}

	@Override
	public <K> OrderedIterable<T> thenBy(Function<T, K> keySelector,
			Comparator<K> comparer) {
		Requires.requireNonNull(keySelector, "The key selector can not be null.");
		Requires.requireNonNull(comparer, "The comparer can not be null.");
		return createOrderedEnumerable(keySelector, comparer, false);
	}

	@Override
	public <K extends Comparable<? super K>> OrderedIterable<T> thenByDescending(Function<T, K> keySelector) {
		return thenByDescending(keySelector, Comparators.<K>Natural());
	}

	@Override
	public <K> OrderedIterable<T> thenByDescending(Function<T, K> keySelector,
			Comparator<K> comparer) {
		Requires.requireNonNull(keySelector, "The key selector can not be null.");
		Requires.requireNonNull(comparer, "The comparer can not be null.");
		return createOrderedEnumerable(keySelector, comparer, true);
	}
}

class OrderedIterator<T, CK> extends UnmodifiableIterator<T> {

	private final Comparator<CK> compositeComparer;
	private final Object[] data, keys;
	private final int[] indexes;
	private final Stack<LeftRight> stack;
	
	@SuppressWarnings("unchecked")
	// Safe cast as data was built from an iterable of T
	public OrderedIterator(Object[] data, Function<T, CK> compositeSelector, Comparator<CK> compositeComparer) {

		this.compositeComparer = compositeComparer;
		this.data = data;
		
		this.indexes = new int[data.length];
		for (int i = 0 ; i < indexes.length ; i++) { 
			indexes[i] = i;
		}
		
		keys = new Object[data.length];
	    for (int i = 0; i < keys.length; i++) 
	    { 
	        keys[i] = compositeSelector.execute((T)data[i]); 
	    } 
	    
	    this.stack = new Stack<>();
	    left = 0;
	    right = data.length - 1;
	    stack.push(new LeftRight(left, right));
	    
	}

	@Override
	public boolean hasNext() {
		return stack.size() > 0 || nextIndex <= right;
	}

	private int nextIndex, left, right;
	@Override
	@SuppressWarnings("unchecked")
	// Safe cast as data is originally an Iterable<T> 
	public T next() {
		if (!hasNext()) {
			throw new NoSuchElementException();
		}
		
		if (right <= left && nextIndex <= right) {
			return (T)data[indexes[nextIndex++]];
		}
		
		do {
			iterateQuickSort();
		} while(nextIndex > right);

		return (T)data[indexes[nextIndex++]];
	}
	
	private void iterateQuickSort() {
		while (true) {
			LeftRight leftRight = stack.pop(); 
			left = leftRight.left; 
			right = leftRight.right; 
			if (right > left) { 
				int pivot = left + (right - left) / 2; 
				int pivotPosition = partition(indexes, keys, left, right, pivot); 
				// Push the right sublist first, so that we *pop* the 
				// left sublist first 
				stack.push(new LeftRight(pivotPosition + 1, right)); 
				stack.push(new LeftRight(left, pivotPosition - 1)); 
			} else {
				return;
			}
		}
	}
	
	// left is the index of the leftmost element of the subarray
	// right is the index of the rightmost element of the subarray (inclusive)
	// number of elements in subarray = right-left+1
	private int partition(int[] indexes, Object[] keys, int left, int right, int pivot) {
		
		int pivotIndex = indexes[pivot];
		Object pivotKey = keys[pivotIndex];
		
		// we move the pivot to the end
		indexes[pivot] = indexes[right]; 
	    indexes[right] = pivotIndex; 
		int storeIndex = left;
		
		for (int i = left ; i < right ; i++) {
			int candidateIndex = indexes[i];
			Object candidateKey = keys[candidateIndex];
			@SuppressWarnings("unchecked")
			// Safe cast as the Composite Keys are the outputs of the compositeSelector<T, CK>
			int comparison = compositeComparer.compare((CK)candidateKey, (CK)pivotKey);
			if (comparison < 0 || (comparison == 0 && candidateIndex < pivotIndex)) { // second condition is for the stable sort
				// Swap storeIndex with the current location 
	            indexes[i] = indexes[storeIndex]; 
	            indexes[storeIndex] = candidateIndex; 
	            storeIndex++; 
			}
		}

		// we move the pivot to its final place
		int tmp = indexes[storeIndex]; 
	    indexes[storeIndex] = indexes[right]; 
	    indexes[right] = tmp; 
		return storeIndex;
	}
	
	private class LeftRight { 
	    private final int left, right; 
	    private LeftRight(int left, int right) { 
	        this.left = left; 
	        this.right = right; 
	    } 
	}
	
}
